(** Typechecking and annotation *)

open Symbols;;
open Environment;;

let last xs = List.hd (List.rev xs);;

(** is_subtype a b tests if a is a subtype of b *)
let rec is_subtype a b = match (a,b) with
    (_, Type_all) -> true                (* source type *)
  | (Type_none, _) -> true               (* sink type *)
  | (Type_integer, Type_integer) -> true
  | (Type_string, Type_string) -> true
  | (Type_arrow xs, Type_arrow ys) -> 
        if List.length xs == List.length ys then     (* xs and ys must be the same length *) 
            let xs' = List.rev xs in
            let ys' = List.rev ys in
            if is_subtype (List.hd xs') (List.hd ys') then   (* last xs must be a subtype of last ys *)
                List.fold_left2 (fun b x y -> b && (is_subtype y x)) true xs' ys'      (* inits xs must be a supertype of inits ys *)
            else
                false
        else
            false
  | _ -> false        (* if it hasn't been recognised, its not a subtype *)
;;

let rec check_expr env e = match e.e_tree with
    Integer _ -> {e_lval = false; e_type = Type_integer; e_tree = e.e_tree; }
  | String _ -> {e_lval = false; e_type = Type_string; e_tree = e.e_tree; }
  | Ident s -> let type_data = get_env env s in
        {e_lval = true; e_type = type_data; e_tree = e.e_tree; }
  | Call (callable, args) -> 
        let callable' = get_env env callable in
        (
        match callable' with
            Type_arrow arg_types -> 
                let args' = List.map (check_expr env) args in
                if List.fold_left2 (fun b x y -> b && (is_subtype x.e_type y)) true (List.rev args') (List.tl(List.rev arg_types)) then
                    { e_lval = false; e_type = List.hd(List.rev arg_types); e_tree = Call (callable, args')}
                else
                    raise (ApyTypeException "Argument Type mismatch")
          | _ -> raise (ApyTypeException "Tried to call a non-callable object")
        )   

and check_stmt env r_type s = match s with
    Assign (e1, e2) -> let e2' = check_expr env e2 in
        let e1' = check_expr env e1 in
        if is_subtype e2'.e_type e1'.e_type then
            Assign (e1', e2')
        else
            raise (ApyTypeException "assignment type mismatch")
  | Stmt_list ss -> Stmt_list (List.map (check_stmt env r_type) ss)
  | Typedef (id, t) -> bind_env env id t; Typedef (id, t)
  | Apodef _ -> s
  | Expr_stmt e -> Expr_stmt (check_expr env e)
  | Funcdef (id, args, body) -> let f_type = get_env env id in
        ( match f_type with Type_arrow es -> if List.length args == (List.length es) -1 then
                    Funcdef (id, args, List.map (check_stmt env (Some (last es))) body)        (** should have some kind of new environment here really *)
                else raise (ApyTypeException "Argument number mismatch")
          | _ -> raise (ApyTypeException "tried to define a function that wasn't one")
        )
  | Return e -> ( 
        match r_type with
            Some t -> let e' = check_expr env e in
               if is_subtype e'.e_type t then Return e'
               else raise (ApyTypeException "Invalid return type")
          | None -> raise (ApyTypeException "Unexpected return statement")
        )
;;
 
